"""Conversation memory management."""

from typing import Dict, Any, Optional, List
from datetime import datetime, timedelta
import logging
import json
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory
from langchain.schema import BaseMessage, HumanMessage, AIMessage
from motor.motor_asyncio import AsyncIOMotorClient
from bson import ObjectId

logger = logging.getLogger(__name__)


class ConversationMemory:
    """Manager for conversation history with MongoDB persistence."""
    
    def __init__(
        self,
        mongodb_client: AsyncIOMotorClient,
        max_memory_size: int = 10000,
        ttl_days: int = 30
    ):
        """Initialize conversation memory manager.
        
        Args:
            mongodb_client: MongoDB client for persistence
            max_memory_size: Maximum characters to keep in memory
            ttl_days: Days to keep conversation history
        """
        self.db = mongodb_client.ai_writing
        self.conversations = self.db.conversations
        self.max_memory_size = max_memory_size
        self.ttl_days = ttl_days
        
        # LangChain memory instances cache
        self._memory_cache: Dict[str, ConversationBufferMemory] = {}
    
    async def create_conversation(
        self,
        user_id: str,
        title: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """Create a new conversation session.
        
        Args:
            user_id: ID of the user
            title: Optional conversation title
            metadata: Optional additional metadata
            
        Returns:
            Created conversation ID
        """
        conversation = {
            "user_id": user_id,
            "title": title or f"Conversation {datetime.utcnow().isoformat()}",
            "messages": [],
            "metadata": metadata or {},
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow(),
            "is_active": True,
            "message_count": 0,
            "total_tokens": 0
        }
        
        try:
            result = await self.conversations.insert_one(conversation)
            conversation_id = str(result.inserted_id)
            
            # Initialize LangChain memory for this conversation
            self._memory_cache[conversation_id] = ConversationBufferMemory(
                return_messages=True,
                memory_key="chat_history"
            )
            
            logger.info(f"Created conversation: {conversation_id}")
            return conversation_id
        except Exception as e:
            logger.error(f"Error creating conversation: {e}")
            raise
    
    async def add_message(
        self,
        conversation_id: str,
        role: str,
        content: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> bool:
        """Add a message to conversation history.
        
        Args:
            conversation_id: ID of the conversation
            role: Message role ('human' or 'ai')
            content: Message content
            metadata: Optional message metadata
            
        Returns:
            True if message added successfully
        """
        message = {
            "role": role,
            "content": content,
            "metadata": metadata or {},
            "timestamp": datetime.utcnow()
        }
        
        try:
            # Update MongoDB
            result = await self.conversations.update_one(
                {"_id": ObjectId(conversation_id)},
                {
                    "$push": {"messages": message},
                    "$set": {"updated_at": datetime.utcnow()},
                    "$inc": {"message_count": 1}
                }
            )
            
            if result.modified_count > 0:
                # Update LangChain memory cache
                if conversation_id in self._memory_cache:
                    if role == "human":
                        self._memory_cache[conversation_id].chat_memory.add_user_message(content)
                    else:
                        self._memory_cache[conversation_id].chat_memory.add_ai_message(content)
                
                # Check and apply memory limits
                await self._apply_memory_limits(conversation_id)
                
                return True
            return False
        except Exception as e:
            logger.error(f"Error adding message: {e}")
            raise
    
    async def get_conversation(
        self,
        conversation_id: str,
        load_messages: bool = True
    ) -> Optional[Dict[str, Any]]:
        """Get a conversation by ID.
        
        Args:
            conversation_id: ID of the conversation
            load_messages: Whether to load message history
            
        Returns:
            Conversation dictionary or None
        """
        try:
            conversation = await self.conversations.find_one(
                {"_id": ObjectId(conversation_id)}
            )
            
            if conversation:
                conversation["_id"] = str(conversation["_id"])
                
                # Load into memory cache if needed
                if load_messages and conversation_id not in self._memory_cache:
                    await self._load_to_memory(conversation_id, conversation["messages"])
                
                return conversation
            return None
        except Exception as e:
            logger.error(f"Error getting conversation: {e}")
            return None
    
    async def get_memory(self, conversation_id: str) -> Optional[ConversationBufferMemory]:
        """Get LangChain memory object for a conversation.
        
        Args:
            conversation_id: ID of the conversation
            
        Returns:
            ConversationBufferMemory instance or None
        """
        # Check cache first
        if conversation_id in self._memory_cache:
            return self._memory_cache[conversation_id]
        
        # Load from database
        conversation = await self.get_conversation(conversation_id, load_messages=True)
        if conversation:
            return self._memory_cache.get(conversation_id)
        
        return None
    
    async def _load_to_memory(
        self,
        conversation_id: str,
        messages: List[Dict[str, Any]]
    ):
        """Load messages into LangChain memory.
        
        Args:
            conversation_id: ID of the conversation
            messages: List of message dictionaries
        """
        memory = ConversationBufferMemory(
            return_messages=True,
            memory_key="chat_history"
        )
        
        for message in messages:
            if message["role"] == "human":
                memory.chat_memory.add_user_message(message["content"])
            else:
                memory.chat_memory.add_ai_message(message["content"])
        
        self._memory_cache[conversation_id] = memory
    
    async def _apply_memory_limits(self, conversation_id: str):
        """Apply memory size limits to a conversation.
        
        Args:
            conversation_id: ID of the conversation
        """
        conversation = await self.conversations.find_one(
            {"_id": ObjectId(conversation_id)},
            {"messages": 1}
        )
        
        if not conversation:
            return
        
        messages = conversation["messages"]
        total_size = sum(len(msg["content"]) for msg in messages)
        
        if total_size > self.max_memory_size:
            # Keep only recent messages that fit within limit
            kept_messages = []
            current_size = 0
            
            for message in reversed(messages):
                msg_size = len(message["content"])
                if current_size + msg_size <= self.max_memory_size:
                    kept_messages.insert(0, message)
                    current_size += msg_size
                else:
                    break
            
            # Update database
            await self.conversations.update_one(
                {"_id": ObjectId(conversation_id)},
                {"$set": {"messages": kept_messages}}
            )
            
            # Reload memory cache
            await self._load_to_memory(conversation_id, kept_messages)
    
    async def clear_conversation(self, conversation_id: str) -> bool:
        """Clear all messages from a conversation.
        
        Args:
            conversation_id: ID of the conversation
            
        Returns:
            True if cleared successfully
        """
        try:
            result = await self.conversations.update_one(
                {"_id": ObjectId(conversation_id)},
                {
                    "$set": {
                        "messages": [],
                        "message_count": 0,
                        "updated_at": datetime.utcnow()
                    }
                }
            )
            
            # Clear memory cache
            if conversation_id in self._memory_cache:
                self._memory_cache[conversation_id].clear()
            
            return result.modified_count > 0
        except Exception as e:
            logger.error(f"Error clearing conversation: {e}")
            raise
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        """Delete a conversation.
        
        Args:
            conversation_id: ID of the conversation
            
        Returns:
            True if deleted successfully
        """
        try:
            result = await self.conversations.delete_one(
                {"_id": ObjectId(conversation_id)}
            )
            
            # Remove from cache
            if conversation_id in self._memory_cache:
                del self._memory_cache[conversation_id]
            
            return result.deleted_count > 0
        except Exception as e:
            logger.error(f"Error deleting conversation: {e}")
            raise
    
    async def list_conversations(
        self,
        user_id: str,
        limit: int = 50,
        skip: int = 0
    ) -> List[Dict[str, Any]]:
        """List conversations for a user.
        
        Args:
            user_id: ID of the user
            limit: Maximum number to return
            skip: Number to skip
            
        Returns:
            List of conversation summaries
        """
        try:
            cursor = self.conversations.find(
                {"user_id": user_id, "is_active": True},
                {
                    "title": 1,
                    "created_at": 1,
                    "updated_at": 1,
                    "message_count": 1
                }
            ).sort("updated_at", -1).skip(skip).limit(limit)
            
            conversations = await cursor.to_list(length=limit)
            for conv in conversations:
                conv["_id"] = str(conv["_id"])
            
            return conversations
        except Exception as e:
            logger.error(f"Error listing conversations: {e}")
            return []
    
    async def cleanup_old_conversations(self):
        """Clean up old conversations based on TTL."""
        cutoff_date = datetime.utcnow() - timedelta(days=self.ttl_days)
        
        try:
            result = await self.conversations.delete_many(
                {"updated_at": {"$lt": cutoff_date}}
            )
            logger.info(f"Cleaned up {result.deleted_count} old conversations")
        except Exception as e:
            logger.error(f"Error cleaning up conversations: {e}")